package com.gree.fms.processor;

import com.google.auto.service.AutoService;
import com.gree.fms.processor.annotations.RocketMQMessageListenerExpand;
import com.gree.fms.processor.annotations.TemplateBean;
import com.gree.fms.processor.constants.ClassRefConstant;
import com.gree.fms.processor.convertors.ExpandToMessageListenerConvertor;
import com.gree.fms.processor.supports.Filters;
import com.sun.source.tree.Tree;
import com.sun.source.util.TreePath;
import com.sun.tools.javac.api.JavacTrees;
import com.sun.tools.javac.code.Flags;
import com.sun.tools.javac.code.Symbol;
import com.sun.tools.javac.code.TypeTag;
import com.sun.tools.javac.processing.JavacProcessingEnvironment;
import com.sun.tools.javac.tree.JCTree;
import com.sun.tools.javac.tree.TreeMaker;
import com.sun.tools.javac.tree.TreeTranslator;
import com.sun.tools.javac.util.Context;
import com.sun.tools.javac.util.ListBuffer;
import com.sun.tools.javac.util.Names;

import javax.annotation.processing.*;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.Element;
import javax.lang.model.element.ElementKind;
import javax.lang.model.element.Name;
import javax.lang.model.element.TypeElement;
import javax.tools.Diagnostic;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Proxy;
import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * @author joyce
 */
@SupportedAnnotationTypes({ClassRefConstant.TEMPLATE_NAME, ClassRefConstant.ANNOTATION_EXPAND_NAME })
@SupportedSourceVersion(SourceVersion.RELEASE_8)
@AutoService(Processor.class)
public class RocketMQMessageListenerExpandProcessor extends AbstractProcessor {
    static String targetAnnotation = "org.apache.rocketmq.spring.annotation.RocketMQMessageListener";

    Messager messager;
    private JavacTrees trees;
    private TreeMaker treeMaker;
    private Names names;

    boolean supportCompiler = true;

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        processingEnv.getMessager().printMessage(Diagnostic.Kind.NOTE, "init:RocketMQMessageListenerExpandProcessor!");
        this.messager = processingEnv.getMessager();
        if (!(processingEnv instanceof JavacProcessingEnvironment)) {
            processingEnv.getMessager().printMessage(Diagnostic.Kind.WARNING, "init:RocketMQMessageListenerExpandProcessor not support by current env!");
            supportCompiler = false;
            try {
                processingEnv = fromProxy(processingEnv);
            } catch (Exception e) {
                processingEnv.getMessager().printMessage(Diagnostic.Kind.WARNING, "tool: not javac class");
                return;
            }
        }
        this.trees = JavacTrees.instance(processingEnv);
        Context context = ((JavacProcessingEnvironment) processingEnv).getContext();
        this.treeMaker = TreeMaker.instance(context);
        this.names = Names.instance(context);
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        if (!supportCompiler) {
            processingEnv.getMessager().printMessage(Diagnostic.Kind.WARNING, "RocketMQMessageListenerExpandProcessor not support by current env!");
            return false;
        }
        processingEnv.getMessager().printMessage(Diagnostic.Kind.NOTE, "RocketMQMessageListenerExpandProcessor calling");
        if (isPresent(targetAnnotation)) {
            Set<? extends Element> templateClass = roundEnv.getElementsAnnotatedWith(TemplateBean.class);
            Set<? extends Element> beanMethod = roundEnv.getElementsAnnotatedWith(RocketMQMessageListenerExpand.class);
            // 前置校验处理 处理
            if (templateClass == null || templateClass.isEmpty() || beanMethod == null || beanMethod.isEmpty()) {
                return false;
            }

            HashMap<ElementKind, List<? extends Element>>  kindGroup = new HashMap<>();
            ListBuffer<Element> classBuffer = new ListBuffer<>();
            classBuffer.addAll(templateClass);
            kindGroup.put(ElementKind.METHOD, new ArrayList<>(beanMethod));
            List<Element> classList = Filters.filters(classBuffer.toList(), trees);
            kindGroup.put(ElementKind.CLASS, classList);

            Map<String, ? extends Element> classNameGroup = classList.stream().collect(Collectors.toMap(Element::toString, Function.identity()));

            HashMap<String, List<JCTree.JCClassDecl>> classMap = new HashMap<>();
            // method 处理
            for (Element methodElement : kindGroup.get(ElementKind.METHOD)) {
                // 处理类名称
                Name methodName = methodElement.getSimpleName();
                // 处理返回对象
                String returnTypeClassName = ((Symbol.MethodSymbol) methodElement).getReturnType().tsym.toString();
                Element targetModifyClass = classNameGroup.get(returnTypeClassName);
                if (targetModifyClass == null) {
                    continue;
                }
                String innerClassName = targetModifyClass.getSimpleName() + "$" + methodName;
                // 修改类信息
                modifyClass(methodElement, targetModifyClass, innerClassName, classMap);
                // 对方法处理
                modifyMethod(methodElement, targetModifyClass, innerClassName, classMap);
            }

            // 统一处理类型信息
            kindGroup.get(ElementKind.CLASS).forEach(c -> {
                // 导包
                addImport(c);
                JCTree tree = trees.getTree(c);
                tree.accept(new TreeTranslator(){

                    @Override
                    public void visitClassDef(JCTree.JCClassDecl jcClassDecl) {
                        super.visitClassDef(jcClassDecl);
                        for (JCTree.JCClassDecl classDecl : classMap.getOrDefault(c.toString(), Collections.emptyList())) {
                            jcClassDecl.defs = jcClassDecl.defs.append(classDecl);
                        }
                    }
                });
            });
        }
        return true;
    }

    // 类是否存在
    private static boolean isPresent(String name) {
        try {
            RocketMQMessageListenerExpandProcessor.class.getClassLoader().loadClass(name);
            return true;
        } catch (ClassNotFoundException e) {
            return false;
        }
    }


    private void modifyClass(Element methodElement, Element targetModifyClass, String innerClassName, HashMap<String, List<JCTree.JCClassDecl>> classMap) {
        RocketMQMessageListenerExpand methodTargetAnnotation = methodElement.getAnnotation(RocketMQMessageListenerExpand.class);
        // 处理返回对象
        String returnTypeClassName = targetModifyClass.toString();

        JCTree tree = trees.getTree(targetModifyClass);
        List<JCTree.JCTypeParameter> jcTypeParameters = new ArrayList<>();
        // 生成子类
        tree.accept(new TreeTranslator(){
            @Override
            public void visitClassDef(JCTree.JCClassDecl jcClassDecl) {
                super.visitClassDef(jcClassDecl);
                // 定义内部类
                JCTree.JCClassDecl innerClass = treeMaker.ClassDef(
                        treeMaker.Modifiers(0),
                        names.fromString(innerClassName),
                        com.sun.tools.javac.util.List.from(jcTypeParameters),
                        treeMaker.Type(((Symbol.ClassSymbol) targetModifyClass).type),
                        com.sun.tools.javac.util.List.nil(),
                        com.sun.tools.javac.util.List.nil()
                );

                // 创建注解
                JCTree.JCAnnotation annotation = treeMaker.Annotation(
                        treeMaker.Ident(names.fromString("RocketMQMessageListener")),
                        ExpandToMessageListenerConvertor.converter(treeMaker, names, methodTargetAnnotation)
                );

                // 将注解添加到内部类的修饰符中
                innerClass.mods = treeMaker.Modifiers(Flags.PUBLIC + Flags.STATIC, com.sun.tools.javac.util.List.of(annotation));
                // 类定义处理
                classMap.computeIfPresent(returnTypeClassName, (k, v) -> {
                    v.add(innerClass);
                    return v;
                });
                classMap.computeIfAbsent(returnTypeClassName, (k) -> {
                    List<JCTree.JCClassDecl> eleList = new ArrayList<>();
                    eleList.add(innerClass);
                    return eleList;
                });
            }

            @Override
            public void visitTypeParameter(JCTree.JCTypeParameter jcTypeParameter) {
                jcTypeParameters.add(jcTypeParameter);
                super.visitTypeParameter(jcTypeParameter);
            }
        });
    }

    private void modifyMethod(Element methodElement, Element targetModifyClass, String innerClassName, HashMap<String, List<JCTree.JCClassDecl>> classMap) {
        JCTree methodTree = trees.getTree(methodElement);
        methodTree.accept(new TreeTranslator(){
            @Override
            public void visitMethodDef(JCTree.JCMethodDecl jcMethodDecl) {
                super.visitMethodDef(jcMethodDecl);
                // 修改方法注解
//                if (classMap.containsKey(jcMethodDecl.sym.owner.toString())) {
//                    // 满足条件则不增强
//                    return;
//                }

                // 修改返回值
                // 先定义 Class 对象
                String newClassPath = targetModifyClass + "$" + innerClassName;
                RocketMQMessageListenerExpand methodTargetAnnotation = methodElement.getAnnotation(RocketMQMessageListenerExpand.class);
                JCTree.JCMethodInvocation initInstance = treeMaker.Apply(
                        com.sun.tools.javac.util.List.nil(),
                        treeMaker.Select(
                                treeMaker.Ident(names.fromString("ClassInstanceConvertor")),
                                names.fromString("instanceInit")
                        ),
                        com.sun.tools.javac.util.List.of(
                                treeMaker.Literal(newClassPath),
                                treeMaker.Literal(methodTargetAnnotation.customize()),
                                treeMaker.Literal(methodTargetAnnotation.tableKey())
                        )
                );

                // 定义反回置类型
                String retTypeClassName = jcMethodDecl.restype instanceof JCTree.JCTypeApply ?
                        ((JCTree.JCTypeApply) jcMethodDecl.restype).clazz.toString() :
                        jcMethodDecl.restype.toString();
                com.sun.tools.javac.util.List<JCTree.JCExpression> classTType = parseRetType(jcMethodDecl);
                JCTree.JCExpression nullType = classTType.isEmpty() ?
                        treeMaker.Ident(names.fromString(retTypeClassName)):
                        treeMaker.TypeApply(treeMaker.Ident(names.fromString(retTypeClassName)), parseRetType(jcMethodDecl));

                JCTree.JCReturn aReturn = treeMaker.Return(treeMaker.TypeCast(nullType, initInstance));

                jcMethodDecl.body.stats = com.sun.tools.javac.util.List.of(aReturn);
                addImport(methodElement);
            }
        });
    }

    private void addImport(Element e, JCTree.JCImport... imports) {
        TreePath treePath = this.trees.getPath(e);
        JCTree.JCCompilationUnit jct = (JCTree.JCCompilationUnit) treePath.getCompilationUnit();
        if (null == jct.defs || jct.defs.isEmpty()) {
            return;
        }

        Set<String> existImport = jct.defs.stream()
                .filter(def -> def.getKind().equals(Tree.Kind.IMPORT))
                .map(JCTree::toString)
                .collect(Collectors.toSet());

        JCTree.JCImport a1 = this.treeMaker.Import(this.memberAccess(ClassRefConstant.CONSUME_MODE), false);
        JCTree.JCImport a2 = this.treeMaker.Import(this.memberAccess(ClassRefConstant.MESSAGE_MODE), false);
        JCTree.JCImport a3 = this.treeMaker.Import(this.memberAccess(ClassRefConstant.SELECTOR_TYPE), false);
        JCTree.JCImport a4 = this.treeMaker.Import(this.memberAccess(ClassRefConstant.RMQ_MESSAGE_LISTENER), false);
        JCTree.JCImport a5 = this.treeMaker.Import(this.memberAccess(ClassRefConstant.INTERFACE_EXPAND_NAME), false);
        JCTree.JCImport a6 = this.treeMaker.Import(this.memberAccess(ClassRefConstant.CONVERTER_NAME), false);

        ListBuffer<JCTree> jcTrees = new ListBuffer<>();
        // 固定
        for (JCTree.JCImport jcImport : com.sun.tools.javac.util.List.of(a1, a2, a3, a4, a5, a6)) {
            if (existImport.contains(jcImport.toString())){
                continue;
            }
            jcTrees.add(jcImport);
        }

        // 动态
        for (JCTree.JCImport jcImport : com.sun.tools.javac.util.List.from(imports)) {
            if (existImport.contains(jcImport.toString())){
                continue;
            }
            jcTrees.add(jcImport);
        }
        // 导入的包要放在最上面， 所以在最后放原始凭证数据
        jcTrees.addAll(jct.defs);
        jct.defs = com.sun.tools.javac.util.List.from(jcTrees);
    }

    private JCTree.JCExpression memberAccess(String components) {
        String[] componentArray = components.split("\\.");
        JCTree.JCExpression expr = this.treeMaker.Ident(names.fromString(componentArray[0]));

        for (int i = 1; i < componentArray.length; ++i) {
            expr = this.treeMaker.Select(expr, names.fromString(componentArray[i]));
        }
        return expr;
    }

    public com.sun.tools.javac.util.List<JCTree.JCExpression> parseRetType(JCTree.JCMethodDecl jcMethodDecl) {
        // 注意 的是这些信息编译后由于类型擦除，此时无任何意义
        ListBuffer<JCTree.JCExpression> nullTypeAndT = new ListBuffer<>();
        com.sun.tools.javac.util.List<JCTree.JCExpression> arguments = jcMethodDecl.restype instanceof JCTree.JCTypeApply ?
                ((JCTree.JCTypeApply) jcMethodDecl.restype).arguments:
                com.sun.tools.javac.util.List.nil();
        if (arguments != null && !arguments.isEmpty()) {
            for (JCTree.JCExpression argument : ((JCTree.JCTypeApply) jcMethodDecl.restype).arguments) {
                JCTree.JCExpression T = treeMaker.Ident(names.fromString(argument.toString()));
                nullTypeAndT.add(T);
            }
        }
        return nullTypeAndT.toList();
    }


    private ProcessingEnvironment fromProxy(Object proxy) throws NoSuchFieldException, IllegalAccessException {
        if (Proxy.isProxyClass(proxy.getClass())) {
            InvocationHandler handler = Proxy.getInvocationHandler(proxy);
            Field hField = handler.getClass().getDeclaredField("h");
            hField.setAccessible(true);
            return (ProcessingEnvironment)hField.get(handler);
        }
        throw new RuntimeException("认命吧使用maven代理编译！");
    }

    public static void main(String[] args) {

    }
}
